#include "chunkmaster.grpc.pb.h"
#include "pool.h"
#include <grpcpp/grpcpp.h>
#include <memory>
#include <vector>
#include <atomic>
#include <condition_variable>
#include <mutex>

extern Pool* g_pool_ptr;

using grpc::Server;
using grpc::ServerBuilder;
using grpc::ServerContext;
using grpc::Status;
using chunkmaster::ChunkMaster;
using chunkmaster::ChunkIdRequest;
using chunkmaster::ChunkIdReply;
using chunkmaster::ChunkIdsRequest;
using chunkmaster::ChunkIdsReply;
using chunkmaster::GetChunkInfoRequest;
using chunkmaster::GetChunkInfoReply;

std::atomic<bool> g_grpc_should_exit{false};
std::condition_variable g_grpc_cv;
std::mutex g_grpc_mutex;

class ChunkMasterServiceImpl final : public ChunkMaster::Service {
public:
    Status AllocateChunkId(ServerContext* context, const ChunkIdRequest* req, ChunkIdReply* reply) override {
        if (!g_pool_ptr) return Status(grpc::StatusCode::UNAVAILABLE, "Pool not ready");
        ChunkMeta meta = g_pool_ptr->alloc_chunkid_balanced();
        reply->set_chunk_id(meta.chunk_id);
        return Status::OK;
    }
    Status AllocateChunkIds(ServerContext* context, const ChunkIdsRequest* req, ChunkIdsReply* reply) override {
        if (!g_pool_ptr) return Status(grpc::StatusCode::UNAVAILABLE, "Pool not ready");
        uint32_t count = req->count();
        for (uint32_t i = 0; i < count; ++i) {
            ChunkMeta meta = g_pool_ptr->alloc_chunkid_balanced();
            reply->add_chunk_ids(meta.chunk_id);
        }
        return Status::OK;
    }
    Status GetChunkInfo(ServerContext* context, const GetChunkInfoRequest* req, GetChunkInfoReply* reply) override {
        if (!g_pool_ptr) return Status(grpc::StatusCode::UNAVAILABLE, "Pool not ready");
        uint64_t chunk_id = req->chunk_id();
        const ChunkMeta* meta = g_pool_ptr->get_chunk_meta(chunk_id);
        if (!meta) {
            return Status(grpc::StatusCode::NOT_FOUND, "Chunk not found");
        }
        
        reply->set_chunk_id(meta->chunk_id);
        reply->set_pg_id(meta->pg_id);
        
        const PG* pg = g_pool_ptr->get_pg(meta->pg_id);
        if (pg) {
            for (const auto& disk_loc : pg->disk_locs) {
                std::string disk_str = "S" + std::to_string(disk_loc.first) + "-D" + std::to_string(disk_loc.second);
                reply->add_disks(disk_str);
            }
        }
        
        return Status::OK;
    }
};

void RunGrpcServer(const std::string& addr) {
    ChunkMasterServiceImpl service;
    ServerBuilder builder;
    builder.AddListeningPort(addr, grpc::InsecureServerCredentials());
    builder.RegisterService(&service);
    std::unique_ptr<Server> server(builder.BuildAndStart());
    std::cout << "gRPC服务已启动: " << addr << std::endl;

    // 等待退出信号
    std::unique_lock<std::mutex> lk(g_grpc_mutex);
    g_grpc_cv.wait(lk, [] { return g_grpc_should_exit.load(); });
    server->Shutdown();
} 